from tkinter import *
from tkinter import filedialog, messagebox, ttk
from text_blind_watermark import TextBlindWatermark

class WatermarkApp:
    def __init__(self, root):
        self.root = root
        self.root.title("文本水印工具")
        self.root.geometry("500x400")
        self.root.resizable(False, False)
        self.root.configure(bg='#f0f0f0')
        
        # 主框架
        main_frame = Frame(root, padx=20, pady=20, bg='#f0f0f0')
        main_frame.pack(fill=BOTH, expand=True)
        
        # 输入区域
        input_frame = Frame(main_frame, bg='#f0f0f0')
        input_frame.pack(fill=X, pady=10)
        
        # 密码输入
        Label(input_frame, text="密码:", bg='#f0f0f0').grid(row=0, column=0, sticky=W, pady=5)
        self.password_entry = Entry(input_frame, show="*", width=30)
        self.password_entry.grid(row=0, column=1, padx=5, pady=5)
        
        # 水印输入
        Label(input_frame, text="水印内容:", bg='#f0f0f0').grid(row=1, column=0, sticky=W, pady=5)
        self.watermark_entry = Entry(input_frame, width=30)
        self.watermark_entry.grid(row=1, column=1, padx=5, pady=5)
        
        # 文件选择
        file_frame = Frame(main_frame, bg='#f0f0f0')
        file_frame.pack(fill=X, pady=10)
        
        Label(file_frame, text="选择文件:", bg='#f0f0f0').grid(row=0, column=0, sticky=W, pady=5)
        self.file_path = StringVar()
        Entry(file_frame, textvariable=self.file_path, state='readonly', width=25).grid(row=0, column=1, padx=5)
        Button(file_frame, text="浏览", command=self.select_file, bg='#007bff', fg='white', activebackground='#0056b3').grid(row=0, column=2, padx=5)
        
        # 操作按钮
        button_frame = Frame(main_frame, bg='#f0f0f0')
        button_frame.pack(fill=X, pady=20)
        
        Button(button_frame, text="添加水印", command=self.add_watermark, 
               bg='#28a745', fg='white', activebackground='#218838', width=15).pack(side=LEFT, padx=10)
        Button(button_frame, text="提取水印", command=self.extract_watermark,
               bg='#17a2b8', fg='white', activebackground='#138496', width=15).pack(side=LEFT, padx=10)
        
        # 状态显示
        status_frame = Frame(main_frame, bg='#f0f0f0')
        status_frame.pack(fill=X, pady=10)
        
        self.status = StringVar()
        self.status.set("准备就绪")
        Label(status_frame, textvariable=self.status, bg='#f0f0f0', fg='#333').pack()
    
    def select_file(self):
        file = filedialog.askopenfilename(filetypes=[("文本文件", "*.txt")])
        if file:
            self.file_path.set(file)
            self.status.set(f"已选择文件: {file}")
    
    def add_watermark(self):
        try:
            password = self.password_entry.get().encode('utf-8')
            watermark = self.watermark_entry.get().encode('utf-8')
            file_path = self.file_path.get()
            
            if not all([password, watermark, file_path]):
                messagebox.showerror("错误", "请填写所有字段并选择文件")
                return
            
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
            
            twm = TextBlindWatermark(pwd=password)
            text_with_wm = twm.add_wm_rnd(text=text, wm=watermark)
            
            save_path = filedialog.asksaveasfilename(defaultextension=".txt", 
                                                    filetypes=[("文本文件", "*.txt")])
            if save_path:
                with open(save_path, 'w', encoding='utf-8') as f:
                    f.write(text_with_wm)
                self.status.set(f"水印添加成功，已保存到: {save_path}")
        except Exception as e:
            messagebox.showerror("错误", f"添加水印失败: {str(e)}")
    
    def extract_watermark(self):
        try:
            password = self.password_entry.get().encode('utf-8')
            file_path = self.file_path.get()
            
            if not all([password, file_path]):
                messagebox.showerror("错误", "请填写密码并选择文件")
                return
            
            with open(file_path, 'r', encoding='utf-8') as f:
                text_with_wm = f.read()
            
            twm = TextBlindWatermark(pwd=password)
            watermark_extract = twm.extract(text_with_wm)
            
            messagebox.showinfo("提取结果", f"提取到的水印内容: {watermark_extract.decode('utf-8')}")
            self.status.set("水印提取成功")
        except Exception as e:
            messagebox.showerror("错误", f"提取水印失败: {str(e)}")

if __name__ == "__main__":
    root = Tk()
    app = WatermarkApp(root)
    root.mainloop()